from __future__ import print_function
import os
from keras.utils.vis_utils import plot_model
import numpy as np
from keras.models import Model
from keras.layers import Input,concatenate, Conv3D,MaxPooling3D,UpSampling3D,PReLU,add,RepeatVector
from keras.optimizers import Adam
from keras.callbacks import ModelCheckpoint
from keras import backend as K
import nibabel as nib
from skimage.transform import resize
from skimage.exposure import rescale_intensity,adjust_gamma,adjust_sigmoid
from skimage.external.tifffile import imshow
from skimage.filters import threshold_local
#hyper param
vol_channel=64
vol_width=128
vol_height=128
smooth = 1.


def augment():
    #unsupported
    return



def load_data(case_num,base_path='/home/biancheng/disk1/testkeras/data/',mode='Train'):
    label_path = sorted([base_path+mode+'/'+ p for p in os.listdir(base_path+mode) if 'Label' in p])
    image_path = sorted([base_path+mode+'/'+ p for p in os.listdir(base_path+mode) if 'Label' not in p])
    case_id=[]
    case_ld=[]
    for i in range(case_num):
        id=nib.load(image_path[i]).get_data()
        ld=nib.load(label_path[i]).get_data()


        id = resize(id, (vol_channel,vol_width, vol_height), preserve_range=True)
        ld = resize(ld, (vol_channel,vol_width, vol_height), preserve_range=True)


        #print(np.unique(ld))
        id=id.reshape((1,) + id.shape)
        ld = ld.reshape((1,) + ld.shape)
        case_id.append(id)
        case_ld.append(ld)
    #train_image = np.array(case_id)
    #train_label = np.array(case_ld)
    return case_id,case_ld


def preprocessing_data_label(datalist,labellist):
    preprocess_img_case=[]
    preprocess_label_case=[]
    for i in range(len(datalist)):#batch
        preprocess_img=[]
        preprocess_label = []
        for j in range(len(datalist[i][0])):#channel
            img = datalist[i][0,j]
            label = labellist[i][0,j]
            img_ = rescale_intensity(img,out_range=np.uint8)/255
            label_ = (label>threshold_local(label,31)).astype('float32')
            print(np.unique(img_))
            preprocess_img.append(img_)#add channel
            preprocess_label.append(label_)


        case_img_ = np.transpose(np.array(preprocess_img),(1,2,0)) # channel width height(case)
        case_label_ = np.transpose(np.array(preprocess_label),(1,2,0)) # channel width height(case)
        #case_img_ = np.array(preprocess_img) # channel width height(case)
        #case_label_ = np.array(preprocess_label) # channel width height(case)
        preprocess_img_case.append(case_img_)#add batch
        preprocess_label_case.append(case_label_)
    case_img = np.array(preprocess_img_case)#batch channel width height
    case_label = np.array(preprocess_label_case)

    return case_img,case_label


def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)


def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)




#output is sorted by(TH Ordering) [batch,depth(channel),width(dim1),height(dim2),channel(dim3)]
def get_vnet():
    inputs = Input((1,vol_width,vol_height,vol_channel))#input is 3d image
    conv_split_todepth16 = concatenate([inputs, inputs, inputs, inputs, inputs, inputs, inputs, inputs, inputs, inputs, inputs, inputs, inputs, inputs,inputs, inputs], axis=1)
    conv_in128_chan16 = Conv3D(filters=16,kernel_size=(5,5,5),strides=(1,1,1),padding='same',data_format='channels_first',name='conv_in128_chan16')(inputs)
    #conv_split_todepth16 = Conv3D(filters=16,kernel_size=(1,1,1),strides=(1,1,1),padding='valid',data_format='channels_first',name='conv_split_todepth16')(inputs)
    addLayer1_1 = add([conv_in128_chan16,conv_split_todepth16])
    outBlock1_1_RELU = PReLU(name='outBlock1_1_RELU')(addLayer1_1)
    #first Block


    down_pooling_128_to_64 = Conv3D(filters=32, kernel_size=(2, 2, 2), strides=(2, 2, 2), padding='valid',data_format='channels_first', name='down_pooling_128_to_64')(outBlock1_1_RELU)
    relu_down_pooling_128_to_64 = PReLU(name='relu_down_pooling_128_to_64')(down_pooling_128_to_64)
    #first pooling

    conv_in64_chan32 = Conv3D(filters=32, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in64_chan32')(relu_down_pooling_128_to_64)
    addLayer1_2 = add([conv_in64_chan32, relu_down_pooling_128_to_64])
    outBlock1_2_RELU = PReLU(name='outBlock1_2_RELU')(addLayer1_2)
    #second block

    down_pooling_64_to_32 = Conv3D(filters=64, kernel_size=(2, 2, 2), strides=(2, 2, 2), padding='valid',data_format='channels_first', name='down_pooling_64_to_32')(outBlock1_2_RELU)
    relu_down_pooling_64_to_32 = PReLU(name='relu_down_pooling_64_to_32')(down_pooling_64_to_32)
    #second pooling

    conv_in32_chan64 = Conv3D(filters=64, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in32_chan64')(relu_down_pooling_64_to_32)
    relu_in32_chan64 = PReLU(name='relu_in32_chan64')(conv_in32_chan64)
    conv_in32_chan64_2 = Conv3D(filters=64, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in32_chan64_2')(relu_in32_chan64)
    addLayer1_3 = add([conv_in32_chan64_2, relu_down_pooling_64_to_32])
    outBlock1_3_RELU = PReLU(name='outBlock1_3_RELU')(addLayer1_3)
    #third block

    down_pooling_32_to_16 = Conv3D(filters=128, kernel_size=(2, 2, 2), strides=(2, 2, 2), padding='valid', data_format='channels_first', name='down_pooling_32_to_16')(outBlock1_3_RELU)
    relu_down_pooling_32_to_16 = PReLU()(down_pooling_32_to_16)
    # third pooling

    conv_in16_chan128 = Conv3D(filters=128, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in16_chan128')(relu_down_pooling_32_to_16)
    relu_in16_chan128 = PReLU(name='relu_in16_chan128')(conv_in16_chan128)
    conv_in16_chan128_2 = Conv3D(filters=128, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in16_chan128_2')(relu_in16_chan128)
    relu_in16_chan128_2 = PReLU(name='relu_in16_chan128_2')(conv_in16_chan128_2)
    conv_in16_chan128_3 = Conv3D(filters=128, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in16_chan128_3')(relu_in16_chan128_2)
    addLayer1_4 = add([conv_in16_chan128_3, relu_down_pooling_32_to_16])
    outBlock1_4_RELU = PReLU(name='outBlock1_4_RELU')(addLayer1_4)
    #forth block

    down_pooling_16_to_8 = Conv3D(filters=256, kernel_size=(2, 2, 2), strides=(2, 2, 2), padding='valid',data_format='channels_first', name='down_pooling_16_to_8')(outBlock1_4_RELU)
    relu_down_pooling_16_to_8 = PReLU(name='relu_down_pooling_16_to_8')(down_pooling_16_to_8)
    # forth pooling


    conv_in8_chan256 = Conv3D(filters=256, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in8_chan256')(relu_down_pooling_16_to_8)
    relu_in8_chan256 = PReLU(name='relu_in8_chan256')(conv_in8_chan256)
    conv_in8_chan256_2 = Conv3D(filters=256, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in8_chan256_2')(relu_in8_chan256)
    relu_in8_chan256_2 = PReLU(name='relu_in8_chan256_2')(conv_in8_chan256_2)
    conv_in8_chan256_3 = Conv3D(filters=256, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in8_chan256_3')(relu_in8_chan256_2)
    addLayer1_5 = add([conv_in8_chan256_3, relu_down_pooling_16_to_8])
    outBlock1_5_RELU = PReLU(name='outBlock1_5_RELU')(addLayer1_5)
    # fifth block

    deconv_in8_chan128 = UpSampling3D(size=(2,2,2),data_format='channels_first',name='deconv_in8_chan128')(outBlock1_5_RELU)
    deconv_in8_chan128 = Conv3D(filters=128, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='valid',data_format='channels_first', name='deconv_in8_chan128_conv')(deconv_in8_chan128)
    relu_deconv_in8_chan128 = PReLU(name='relu_deconv_in8_chan128')(deconv_in8_chan128)
    concat_in16_concat = concatenate([relu_deconv_in8_chan128,outBlock1_4_RELU],axis=1)
    # first depool


    conv_in16_chan128_right = Conv3D(filters=256, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in16_chan128_right')(concat_in16_concat)
    relu_conv_in16_chan128_right = PReLU(name='relu_conv_in16_chan128_right')(conv_in16_chan128_right)
    conv_in16_chan128_2_right = Conv3D(filters=256, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in16_chan128_2_right')(relu_conv_in16_chan128_right)
    relu_in16_chan128_2_right = PReLU(name='relu_in16_chan128_2_right')(conv_in16_chan128_2_right)
    conv_in16_chan128_3_right = Conv3D(filters=256, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in16_chan128_3_right')(relu_in16_chan128_2_right)
    addLayer2_4 = add([conv_in16_chan128_3_right, concat_in16_concat])
    outBlock2_4_RELU = PReLU(name='outBlock2_4_RELU')(addLayer2_4)
    #sixth block


    deconv_in16_chan64 = UpSampling3D(size=(2, 2, 2), data_format='channels_first', name='deconv_in16_chan64')(outBlock2_4_RELU)
    deconv_in16_chan64 = Conv3D(filters=64, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='valid',data_format='channels_first', name='deconv_in16_chan64_conv')(deconv_in16_chan64)
    relu_deconv_in16_chan64 = PReLU(name='relu_deconv_in16_chan64')(deconv_in16_chan64)
    concat_in32_concat = concatenate([relu_deconv_in16_chan64, outBlock1_3_RELU], axis=1)
    # second depool


    conv_in32_chan64_right = Conv3D(filters=128, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in32_chan64_right')(concat_in32_concat)
    relu_conv_in32_chan64_right = PReLU()(conv_in32_chan64_right)
    conv_in32_chan64_2_right = Conv3D(filters=128, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in32_chan64_2_right')(relu_conv_in32_chan64_right)
    addLayer2_3 = add([conv_in32_chan64_2_right, concat_in32_concat])
    outBlock2_3_RELU = PReLU(name='outBlock2_3_RELU')(addLayer2_3)
    # seventh block




    deconv_in32_chan32 = UpSampling3D(size=(2, 2, 2), data_format='channels_first', name='deconv_in32_chan32')(outBlock2_3_RELU)
    deconv_in32_chan32 = Conv3D(filters=32, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='valid',data_format='channels_first', name='deconv_in32_chan32_conv')(deconv_in32_chan32)
    relu_deconv_in32_chan32 = PReLU(name='relu_deconv_in32_chan32')(deconv_in32_chan32)
    concat_in64_concat = concatenate([relu_deconv_in32_chan32, outBlock1_2_RELU], axis=1)
    # thrid depool


    conv_in64_chan32_right = Conv3D(filters=64, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same',data_format='channels_first', name='conv_in64_chan32_right')(concat_in64_concat)
    addLayer2_2 = add([conv_in64_chan32_right, concat_in64_concat])
    outBlock2_2_RELU = PReLU(name='outBlock2_2_RELU')(addLayer2_2)
    # eighth block


    deconv_in64_chan16 = UpSampling3D(size=(2, 2, 2), data_format='channels_first', name='deconv_in64_chan16')(outBlock2_2_RELU)
    deconv_in64_chan16 = Conv3D(filters=16, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='valid',data_format='channels_first', name='deconv_in64_chan16_conv')(deconv_in64_chan16)
    relu_deconv_in64_chan16 = PReLU(name='relu_deconv_in64_chan16')(deconv_in64_chan16)
    concat_in128_concat = concatenate([relu_deconv_in64_chan16, outBlock1_1_RELU], axis=1)
    # forth depool


    conv_in128_chan16_right = Conv3D(filters=32, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same', data_format='channels_first', name='conv_in128_chan16_right')(concat_in128_concat)
    addLayer2_1 = add([conv_in128_chan16_right, concat_in128_concat])
    outBlock2_1_RELU = PReLU(name='outBlock2_1_RELU')(addLayer2_1)
    # night block


    conv_in128_chan2_right =  Conv3D(filters=2, kernel_size=(5, 5, 5), strides=(1, 1, 1), padding='same', data_format='channels_first', name='conv_in128_chan2_right')(outBlock2_1_RELU)
    relu_conv_in128_chan2_right = PReLU(name='relu_conv_in128_chan2_right')(conv_in128_chan2_right)
    conv_in128_chan2_right_2 = Conv3D(filters=2, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='valid',data_format='channels_first', name='conv_in128_chan2_right_2')(relu_conv_in128_chan2_right)


    model = Model(inputs=[inputs], outputs=[conv_in128_chan2_right_2])

    model.compile(optimizer=Adam(lr=1e-5), loss=dice_coef_loss, metrics=[dice_coef])

    return model

train_image,train_label=load_data(mode='Train',case_num=15)


model = get_vnet()
image_trainset,label_trainset=preprocessing_data_label(train_image,train_label)

#fig=plt.figure(figsize=(20,4))
#imshow(label_set,figure=fig)
#plt.show()


image_trainset = image_trainset.reshape(image_trainset.shape[0:1]+(1,)+image_trainset.shape[1:])
label_trainset = image_trainset.reshape(label_trainset.shape[0:1]+(1,)+label_trainset.shape[1:])


#draw V-net model
plot_model(model,to_file='V-net.png')




model.summary()
#model.fit(image_trainset, label_trainset, batch_size=1, nb_epoch=200, verbose=1,validation_split=0.1)#validation_data=(image_testset,label_testset))
